#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <assert.h>
#include <mpi.h>
#include <gsl/gsl_math.h>

#include "halofinder.h"

//from gadget-2 peano.c functions
typedef union {
  float d;
  peanokey key;
} floatKeyUnion;

typedef struct {
  int index;
  floatKeyUnion u;
} SortPart;

static int compPeanoKey(const void *a, const void *b)
{
  if(((const SortPart*)a)->u.key > ((const SortPart*)b)->u.key)
    return 1;
  else if(((const SortPart*)a)->u.key < ((const SortPart*)b)->u.key)
    return -1;
  else
    return 0;
}

static int quadrants[24][2][2][2] = {
  /* rotx=0, roty=0-3 */
  {{{0, 7}, {1, 6}}, {{3, 4}, {2, 5}}},
  {{{7, 4}, {6, 5}}, {{0, 3}, {1, 2}}},
  {{{4, 3}, {5, 2}}, {{7, 0}, {6, 1}}},
  {{{3, 0}, {2, 1}}, {{4, 7}, {5, 6}}},
  /* rotx=1, roty=0-3 */
  {{{1, 0}, {6, 7}}, {{2, 3}, {5, 4}}},
  {{{0, 3}, {7, 4}}, {{1, 2}, {6, 5}}},
  {{{3, 2}, {4, 5}}, {{0, 1}, {7, 6}}},
  {{{2, 1}, {5, 6}}, {{3, 0}, {4, 7}}},
  /* rotx=2, roty=0-3 */
  {{{6, 1}, {7, 0}}, {{5, 2}, {4, 3}}},
  {{{1, 2}, {0, 3}}, {{6, 5}, {7, 4}}},
  {{{2, 5}, {3, 4}}, {{1, 6}, {0, 7}}},
  {{{5, 6}, {4, 7}}, {{2, 1}, {3, 0}}},
  /* rotx=3, roty=0-3 */
  {{{7, 6}, {0, 1}}, {{4, 5}, {3, 2}}},
  {{{6, 5}, {1, 2}}, {{7, 4}, {0, 3}}},
  {{{5, 4}, {2, 3}}, {{6, 7}, {1, 0}}},
  {{{4, 7}, {3, 0}}, {{5, 6}, {2, 1}}},
  /* rotx=4, roty=0-3 */
  {{{6, 7}, {5, 4}}, {{1, 0}, {2, 3}}},
  {{{7, 0}, {4, 3}}, {{6, 1}, {5, 2}}},
  {{{0, 1}, {3, 2}}, {{7, 6}, {4, 5}}},
  {{{1, 6}, {2, 5}}, {{0, 7}, {3, 4}}},
  /* rotx=5, roty=0-3 */
  {{{2, 3}, {1, 0}}, {{5, 4}, {6, 7}}},
  {{{3, 4}, {0, 7}}, {{2, 5}, {1, 6}}},
  {{{4, 5}, {7, 6}}, {{3, 2}, {0, 1}}},
  {{{5, 2}, {6, 1}}, {{4, 3}, {7, 0}}}
};


static int rotxmap_table[24] = { 4, 5, 6, 7, 8, 9, 10, 11,
				 12, 13, 14, 15, 0, 1, 2, 3, 17, 18, 19, 16, 23, 20, 21, 22
};

static int rotymap_table[24] = { 1, 2, 3, 0, 16, 17, 18, 19,
				 11, 8, 9, 10, 22, 23, 20, 21, 14, 15, 12, 13, 4, 5, 6, 7
};

static int rotx_table[8] = { 3, 0, 0, 2, 2, 0, 0, 1 };
static int roty_table[8] = { 0, 1, 1, 2, 2, 3, 3, 0 };

static int sense_table[8] = { -1, -1, -1, +1, +1, -1, -1, -1 };

/* vars needed for inverse peano-hilbert construction
  
  static int flag_quadrants_inverse = 1;
  static char quadrants_inverse_x[24][8];
  static char quadrants_inverse_y[24][8];
  static char quadrants_inverse_z[24][8];
*/

/*! This function computes a Peano-Hilbert key for an integer triplet (x,y,z),
 *  with x,y,z in the range between 0 and 2^bits-1.
 */
peanokey peano_hilbert_key(int x, int y, int z, int bits)
{
  int i, quad, bitx, bity, bitz;
  int mask, rotation, rotx, roty, sense;
  peanokey key;
  
  mask = 1 << (bits - 1);
  key = 0;
  rotation = 0;
  sense = 1;
  
  for(i = 0; i < bits; i++, mask >>= 1)
    {
      bitx = (x & mask) ? 1 : 0;
      bity = (y & mask) ? 1 : 0;
      bitz = (z & mask) ? 1 : 0;

      quad = quadrants[rotation][bitx][bity][bitz];

      key <<= 3;
      key += (sense == 1) ? (quad) : (7 - quad);

      rotx = rotx_table[quad];
      roty = roty_table[quad];
      sense *= sense_table[quad];

      while(rotx > 0)
	{
	  rotation = rotxmap_table[rotation];
	  rotx--;
	}

      while(roty > 0)
	{
	  rotation = rotymap_table[rotation];
	  roty--;
	}
    }

  return key;
}

void build_parts_kdtree(void)
{
  PartskdTreeData = buildkdTree(PartsPx,PartsPy,PartsPz,(int) NumParts,NULL);
}

void destroy_parts_kdtree(void)
{
  destroykdTree(PartskdTreeData);
  PartskdTreeData = NULL;
}

void reorder_parts_tree(void)
{
  int i,j;
  float domainLengths[3];
  float *pos[6];
  SortPart *sp;
  double sortTime = 0.0;
  int N,x,y,z;
  double dx,dy,dz;
  
#ifdef DEBUG
  if(ThisTask == 0)
    fprintf(stderr,"%d: redordering parts according to peano-hilbert index\n",ThisTask);
#endif

  //setup
  pos[0] = PartsPx;
  pos[1] = PartsPy;
  pos[2] = PartsPz;
  pos[3] = PartsVx;
  pos[4] = PartsVy;
  pos[5] = PartsVz;
  
  domainLengths[0] = (float) (haloFinderData.BoxLengthSnapshot/((float) (NumDomains[0])) + 2.0*haloFinderData.DomainBuffSize);
  domainLengths[1] = (float) (haloFinderData.BoxLengthSnapshot/((float) (NumDomains[1])) + 2.0*haloFinderData.DomainBuffSize);
  domainLengths[2] = (float) (haloFinderData.BoxLengthSnapshot/((float) (NumDomains[2])) + 2.0*haloFinderData.DomainBuffSize);
  
  N = 1;
  N = N << BITS_PER_DIMENSION;
  
  dx = domainLengths[0]/((double) (N));
  dy = domainLengths[1]/((double) (N));
  dz = domainLengths[2]/((double) (N));
  
  for(i=0;i<NumParts;++i)
    {
      if(!(PartsPx[i] >= 0.0 && PartsPx[i] <= domainLengths[0]))
	fprintf(stderr,"%d: part %d - px,py,pz = %f|%f|%f vx,vy,vz = %f|%f|%f\n",ThisTask,i,
		PartsPx[i],PartsPy[i],PartsPz[i],PartsVx[i],PartsVy[i],PartsVz[i]);
      assert(PartsPx[i] >= 0.0 && PartsPx[i] <= domainLengths[0]);
      
      if(!(PartsPy[i] >= 0.0 && PartsPy[i] <= domainLengths[1]))
	fprintf(stderr,"%d: part %d - px,py,pz = %f|%f|%f vx,vy,vz = %f|%f|%f\n",ThisTask,i,
		PartsPx[i],PartsPy[i],PartsPz[i],PartsVx[i],PartsVy[i],PartsVz[i]);
      assert(PartsPy[i] >= 0.0 && PartsPy[i] <= domainLengths[1]);
      
      if(!(PartsPz[i] >= 0.0 && PartsPz[i] <= domainLengths[2]))
	fprintf(stderr,"%d: part %d - px,py,pz = %f|%f|%f vx,vy,vz = %f|%f|%f\n",ThisTask,i,
		PartsPx[i],PartsPy[i],PartsPz[i],PartsVx[i],PartsVy[i],PartsVz[i]);
      assert(PartsPz[i] >= 0.0 && PartsPz[i] <= domainLengths[2]);
    }
  
  PartsInDomainMark = (short int*)malloc(sizeof(short int)*NumParts);
  assert(PartsInDomainMark != NULL);
  for(j=0;j<NumPartsInDomain;++j)
    PartsInDomainMark[j] = 1;
  for(j=NumPartsInDomain;j<NumParts;++j)
    PartsInDomainMark[j] = 0;
  
#ifdef DEBUG
  if(ThisTask == 0)
    fprintf(stderr,"%d: starting sort of parts\n",ThisTask);
#endif
  sp = (SortPart*)malloc(sizeof(SortPart)*NumParts);
  assert(sp != NULL);
#ifdef USEOPENMP
#pragma omp parallel for schedule(static) default(none) private(i,x,y,z,N) shared(NumParts,dx,dy,dz,PartsPx,PartsPy,PartsPz,sp)
#endif
  for(i=0;i<NumParts;++i)
    {
      sp[i].index = i;
      x = (int) (PartsPx[i]/dx);
      if(x == N)
	--x;
      y = (int) (PartsPy[i]/dy);
      if(y == N)
	--y;
      z = (int) (PartsPz[i]/dz);
      if(z == N)
	--z;
      sp[i].u.key = peano_hilbert_key(x,y,z,BITS_PER_DIMENSION);
    }
  sortTime -= MPI_Wtime();
  qsort(sp,(size_t) NumParts,sizeof(SortPart),compPeanoKey);
  sortTime += MPI_Wtime();
#ifdef DEBUG
  if(ThisTask == 0)  
    fprintf(stderr,"%d: done with sort of parts, time = %f [s]\n",ThisTask,sortTime);
#endif

  for(i=0;i<6;++i)
    {
#ifdef DEBUG
      if(ThisTask == 0)
	fprintf(stderr,"%d: on pos = %d\n",ThisTask,i);
#endif
      
      for(j=0;j<NumParts;++j)
	sp[j].u.d = pos[i][sp[j].index];
      
      for(j=0;j<NumParts;++j)
	pos[i][j] = sp[j].u.d;
    }
  
#ifdef DEBUG
  if(ThisTask == 0)
    fprintf(stderr,"%d: on parts in domain mark\n",ThisTask);
#endif
  for(j=0;j<NumParts;++j)
    sp[j].u.d = (float) (PartsInDomainMark[sp[j].index]);
  for(j=0;j<NumParts;++j)
    PartsInDomainMark[j] = (short int) (sp[j].u.d);
  
  free(sp);
    
#ifdef DEBUG
  if(ThisTask == 0)
    fprintf(stderr,"%d: done redordering parts according to peano-hilbert index\n",ThisTask);
#endif
}

void destroyhalos(void)
{
  if(NumHalos > 0)
    {
      free(Halos);
      NumHalos = 0;
      Halos = NULL;
    }
}

void reallocparts(long NumExtra)
{
  float *tmp;
  
  if(NumExtra != 0)
    {
      tmp = (float*)realloc(PartsPx,sizeof(float)*(NumParts + NumExtra));
      assert(tmp != NULL);
      PartsPx = tmp;
      
      tmp = (float*)realloc(PartsPy,sizeof(float)*(NumParts + NumExtra));
      assert(tmp != NULL);
      PartsPy = tmp;
      
      tmp = (float*)realloc(PartsPz,sizeof(float)*(NumParts + NumExtra));
      assert(tmp != NULL);
      PartsPz = tmp;
      
      tmp = (float*)realloc(PartsVx,sizeof(float)*(NumParts + NumExtra));
      assert(tmp != NULL);
      PartsVx = tmp;
      
      tmp = (float*)realloc(PartsVy,sizeof(float)*(NumParts + NumExtra));
      assert(tmp != NULL);
      PartsVy = tmp;
      
      tmp = (float*)realloc(PartsVz,sizeof(float)*(NumParts + NumExtra));
      assert(tmp != NULL);
      PartsVz = tmp;
      
      NumParts += NumExtra;
    }
}

void destroyparts(void)
{
  if(NumParts > 0)
    {
      NumParts = 0;
      NumPartsInDomain = 0;
      
      free(PartsPx);
      PartsPx = NULL;
      free(PartsPy);
      PartsPy = NULL;
      free(PartsPz);
      PartsPz = NULL;
      
      free(PartsVx);
      PartsVx = NULL;
      free(PartsVy);
      PartsVy = NULL;
      free(PartsVz);
      PartsVz = NULL;
      
      if(PartsInDomainMark != NULL)
	{
	  free(PartsInDomainMark);
	  PartsInDomainMark = NULL;
	}
    }
}

typedef struct {
  long ind;
  float val;
} SortHaloVal;

static int compSortHaloVal(const void *a, const void *b)
{
  if(((const SortHaloVal*)a)->val > ((const SortHaloVal*)b)->val)
    return -1;
  else if(((const SortHaloVal*)a)->val < ((const SortHaloVal*)b)->val)
    return 1;
  else
    return 0;
}

static void sort_halos_val(SortHaloVal *sh)
{
  long i,*rank;
  long rankSource,dest,rankSave;
  Halo haloSave,haloSource;
  
  qsort(sh,(size_t) NumHalos,sizeof(SortHaloVal),compSortHaloVal);
  
  rank = (long*)malloc(sizeof(long)*NumHalos);
  assert(rank != NULL);
 for(i=0;i<NumHalos;++i)
    rank[sh[i].ind] = i;
  
  free(sh);
  
  for(i=0;i<NumHalos;++i) /* reoder with an in-place algorithm - see Gadget-2 for details - destroys rank */
    {
      if(i != rank[i])
	{
	  haloSource = Halos[i];
	  rankSource = rank[i];
	  dest = rank[i];
                          
	  do
	    {
	      haloSave = Halos[dest];
	      rankSave = rank[dest];
                              
	      Halos[dest] = haloSource;
	      rank[dest] = rankSource;
              
	      if(dest == i)
		break;
	      
	      haloSource = haloSave;
	      rankSource = rankSave;
              
	      dest = rankSource;
	    }
	  while(1);
	}
    }
  free(rank);
}

void sort_halos_mass(void)
{
  long i;
  SortHaloVal *sh;
  
  sh = (SortHaloVal*)malloc(sizeof(SortHaloVal)*NumHalos);

#ifdef USEOPENMP
#pragma omp parallel for schedule(static) default(none) private(i) shared(NumHalos,Halos,sh)
#endif
  for(i=0;i<NumHalos;++i)
    {
      sh[i].ind = i;
      sh[i].val = Halos[i].mass;
    }
  
  sort_halos_val(sh);
}

void sort_halos_vcirc(void)
{
  long i;
  SortHaloVal *sh;
  
  sh = (SortHaloVal*)malloc(sizeof(SortHaloVal)*NumHalos);

#ifdef USEOPENMP
#pragma omp parallel for schedule(static) default(none) private(i) shared(NumHalos,Halos,sh)
#endif
  for(i=0;i<NumHalos;++i)
    {
      sh[i].ind = i;
      sh[i].val = Halos[i].vcirc;
    }
  
  sort_halos_val(sh);
}

void wrapdist(float *dx, float rcube, float rhalf)
{
  float dxmag = (float) (fabs(*dx));
  
  if(dxmag > rhalf)
    {
      dxmag = rcube - dxmag;
      
      if(*dx > 0)
	dxmag = -dxmag;
    }
  *dx = dxmag;
}

void wrappos(float *dx, float rbox)
{
  long i;
  
  i = 0;
  while(*dx < 0.0)
    {
      *dx = *dx + rbox;
      ++i;
      
      if(i >= 100)
	{
	  fprintf(stderr,"dx = %f, rbox = %f\n",*dx,rbox);
	  assert(0);
	}
    }
  
  i = 0;
  while(*dx > rbox)
    {
      *dx = *dx - rbox;
      ++i;
      
      if(i >= 100)
	{
	  fprintf(stderr,"dx = %f, rbox = %f\n",*dx,rbox);
	  assert(0);
	}
    }
}
